#include "NHook.Native.h"
#include "NHook.Native.cpp"

BeginNHook

#define NT_SUCCESS(x) ((x) >= 0)
#define STATUS_INFO_LENGTH_MISMATCH 0xc0000004

#define SystemHandleInformation 16
#define ObjectBasicInformation 0
#define ObjectNameInformation 1
#define ObjectTypeInformation 2

	typedef int (NTAPI *_NtQuerySystemInformation)(
	ULONG SystemInformationClass,
	PVOID SystemInformation,
	ULONG SystemInformationLength,
	PULONG ReturnLength
	);
typedef int (NTAPI *_NtDuplicateObject)(
	HANDLE SourceProcessHandle,
	HANDLE SourceHandle,
	HANDLE TargetProcessHandle,
	PHANDLE TargetHandle,
	ACCESS_MASK DesiredAccess,
	ULONG Attributes,
	ULONG Options
	);
typedef int (NTAPI *_NtQueryObject)(
	HANDLE ObjectHandle,
	ULONG ObjectInformationClass,
	PVOID ObjectInformation,
	ULONG ObjectInformationLength,
	PULONG ReturnLength
	);

typedef struct _UNICODE_STRING
{
	USHORT Length;
	USHORT MaximumLength;
	PWSTR Buffer;
} UNICODE_STRING, *PUNICODE_STRING;

typedef struct _SYSTEM_HANDLE
{
	ULONG ProcessId;
	BYTE ObjectTypeNumber;
	BYTE Flags;
	USHORT Handle;
	PVOID Object;
	ACCESS_MASK GrantedAccess;
} SYSTEM_HANDLE, *PSYSTEM_HANDLE;

typedef struct _SYSTEM_HANDLE_INFORMATION
{
	ULONG HandleCount;
	SYSTEM_HANDLE Handles[1];
} SYSTEM_HANDLE_INFORMATION, *PSYSTEM_HANDLE_INFORMATION;

typedef enum _POOL_TYPE
{
	NonPagedPool,
	PagedPool,
	NonPagedPoolMustSucceed,
	DontUseThisType,
	NonPagedPoolCacheAligned,
	PagedPoolCacheAligned,
	NonPagedPoolCacheAlignedMustS
} POOL_TYPE, *PPOOL_TYPE;

typedef struct _OBJECT_TYPE_INFORMATION
{
	UNICODE_STRING Name;
	ULONG TotalNumberOfObjects;
	ULONG TotalNumberOfHandles;
	ULONG TotalPagedPoolUsage;
	ULONG TotalNonPagedPoolUsage;
	ULONG TotalNamePoolUsage;
	ULONG TotalHandleTableUsage;
	ULONG HighWaterNumberOfObjects;
	ULONG HighWaterNumberOfHandles;
	ULONG HighWaterPagedPoolUsage;
	ULONG HighWaterNonPagedPoolUsage;
	ULONG HighWaterNamePoolUsage;
	ULONG HighWaterHandleTableUsage;
	ULONG InvalidAttributes;
	GENERIC_MAPPING GenericMapping;
	ULONG ValidAccess;
	BOOLEAN SecurityRequired;
	BOOLEAN MaintainHandleCount;
	USHORT MaintainTypeList;
	POOL_TYPE PoolType;
	ULONG PagedPoolUsage;
	ULONG NonPagedPoolUsage;
} OBJECT_TYPE_INFORMATION, *POBJECT_TYPE_INFORMATION;

PVOID GetLibraryProcAddress(PSTR LibraryName, PSTR ProcName)
{
	return GetProcAddress(GetModuleHandleA(LibraryName), ProcName);
}


public enum class HandleType
{
	File,
	Key,
	Directory,
	Other
};

public ref class ObjectHandle
{
public : static List<ObjectHandle^>^ GetObjectHandles(ProcessInformation^ processInfo)
		 {
			 List<ObjectHandle^>^ result = gcnew List<ObjectHandle^>();

			 _NtQuerySystemInformation NtQuerySystemInformation = (_NtQuerySystemInformation)GetLibraryProcAddress("ntdll.dll", "NtQuerySystemInformation");
			 _NtDuplicateObject NtDuplicateObject =	 (_NtDuplicateObject)GetLibraryProcAddress("ntdll.dll", "NtDuplicateObject");
			 _NtQueryObject NtQueryObject =  (_NtQueryObject)GetLibraryProcAddress("ntdll.dll", "NtQueryObject");

			 auto pHandle = processInfo->ProcessHandle->DangerousGetHandle().ToPointer();
			 ULONG needed;
			 NtQuerySystemInformation(
				 SystemHandleInformation,
				 NULL,
				 0,
				 &needed
				 );

			 int status;

			 auto handleInfoSize = 0x1000;
			 auto handleInfoUnsafe = (PSYSTEM_HANDLE_INFORMATION)malloc(handleInfoSize);

			 /* NtQuerySystemInformation won't give us the correct buffer size, 
			 so we guess by doubling the buffer size. */
			 while ((status = NtQuerySystemInformation(
				 SystemHandleInformation,
				 handleInfoUnsafe,
				 handleInfoSize,
				 NULL
				 )) == STATUS_INFO_LENGTH_MISMATCH)
				 handleInfoUnsafe = (PSYSTEM_HANDLE_INFORMATION)realloc(handleInfoUnsafe, handleInfoSize *= 2);

			 auto_ptr<SYSTEM_HANDLE_INFORMATION> handleInfo(handleInfoUnsafe);
			 for (UInt32 i = 0; i < handleInfo->HandleCount; i++)
			 {
				 SYSTEM_HANDLE handle = handleInfo->Handles[i];
				 if(handle.ProcessId != processInfo->ProcessId)
					 continue;
				 HANDLE dupHandle = NULL;
				 /* Duplicate the handle so we can query it. */
				 if (!NT_SUCCESS(NtDuplicateObject(
					 pHandle,
					 (HANDLE)handle.Handle,
					 GetCurrentProcess(),
					 &dupHandle,
					 0,
					 0,
					 0
					 )))
				 {
					 continue;
				 }
				 SafeHandle safeDup(IntPtr(dupHandle),true);


				 NtQueryObject(
					 dupHandle,
					 ObjectTypeInformation,
					 NULL,
					 0,
					 &needed);

				 auto_ptr<OBJECT_TYPE_INFORMATION> objectTypeInfo((POBJECT_TYPE_INFORMATION)malloc(needed));
				 if (!NT_SUCCESS(NtQueryObject(
					 dupHandle,
					 ObjectTypeInformation,
					 objectTypeInfo.get(),
					 needed,
					 &needed
					 )))
				 {
					 result->Add(gcnew ObjectHandle(&handle,NULL,NULL));
					 continue;
				 }

				 /* Query the object name (unless it has an access of 
				 0x0012019f, on which NtQueryObject could hang. */
				 if (handle.GrantedAccess == 0x0012019f)
				 {
					 result->Add(gcnew ObjectHandle(&handle,objectTypeInfo.get(),NULL));
					 continue;
				 }

				 NtQueryObject(
					 dupHandle,
					 ObjectNameInformation,
					 NULL,
					 0,
					 &needed
					 );

				 auto_ptr<UNICODE_STRING> objectNameInfo((PUNICODE_STRING)malloc(needed));
				 NtQueryObject(
					 dupHandle,
					 ObjectNameInformation,
					 objectNameInfo.get(),
					 needed,
					 &needed
					 );


				 result->Add(gcnew ObjectHandle(&handle, objectTypeInfo.get(), objectNameInfo.get()));

			 }
			 return result;
		 }


private: ObjectHandle(SYSTEM_HANDLE* handle, POBJECT_TYPE_INFORMATION objectTypeInfo, PUNICODE_STRING name)
		 {
			 if(objectTypeInfo != NULL)
			 {
				 AffectHandleTypeEnum(objectTypeInfo);
				 _TypeString = UnicodeToCSString(&objectTypeInfo->Name);
			 }
			 if(name != NULL)
			 {
				 _Name = UnicodeToCSString(name);
			 }
			 _Id = handle->Handle;
		 }

private: int _Id;
public: property int Id
		{
			int get()
			{
				return _Id;
			}
		}
public: void AffectHandleTypeEnum(POBJECT_TYPE_INFORMATION objectTypeInfo)
		{
			_Type = AffectIfEquals(HandleType::File, "File", &objectTypeInfo->Name);
			if(_Type != HandleType::Other)
				return;
			_Type = AffectIfEquals(HandleType::Key, "Key", &objectTypeInfo->Name);
			if(_Type != HandleType::Other)
				return;
			_Type = AffectIfEquals(HandleType::Directory, "Directory", &objectTypeInfo->Name);
		}

private: String^ UnicodeToCSString(PUNICODE_STRING unicodeStr)
		 {
			 if(unicodeStr->Length == 0)
				 return nullptr;
			 auto csArray = Util::ToCSharpArray((BYTE*)unicodeStr->Buffer, unicodeStr->Length);
			 return System::Text::Encoding::Unicode->GetString(csArray,0,csArray->Length);
		 }
private: HandleType AffectIfEquals(HandleType type, String^ strType, PUNICODE_STRING unicode)
		 {
			 if(strType == UnicodeToCSString(unicode))
				 return type;
			 else
				 return HandleType::Other;
		 }

private: HandleType _Type;
public: property HandleType Type
		{
			HandleType get()
			{
				return _Type;
			}
			void set(HandleType value)
			{
				_Type = value;
			}
		}
private: String^ _TypeString;
public: property String^ TypeString
		{
			String^ get()
			{
				return _TypeString;
			}
			void set(String^ value)
			{
				_TypeString = value;
			}
		}

private: String^ _Name;
public: property String^ Name
		{
			String^ get()
			{
				return _Name;
			}
			void set(String^ value)
			{
				_Name = value;
			}
		}

		private: String^ ValueOrEmpty(String^ value)
				 {
					 if(value == nullptr)
						 return String::Empty;
					 return value;
				 }
public: virtual String^ ToString() override
		{
			return Id.ToString() + " : " + ValueOrEmpty(TypeString) + "[" +ValueOrEmpty(Name)+ "]";
		}
};


EndNHook